package com.dynamic.mybatis.core.session;

import com.baomidou.mybatisplus.core.toolkit.Assert;
import com.baomidou.mybatisplus.core.toolkit.CollectionUtils;
import com.baomidou.mybatisplus.core.toolkit.ExceptionUtils;
import com.dynamic.mybatis.core.metadata.DynamicMappedStatement;
import org.apache.ibatis.executor.BatchResult;
import org.apache.ibatis.reflection.ExceptionUtil;
import org.apache.ibatis.session.ExecutorType;
import org.apache.ibatis.session.SqlSession;
import org.springframework.transaction.annotation.Transactional;
import org.springframework.transaction.support.TransactionSynchronizationManager;

import java.util.Arrays;
import java.util.Collection;
import java.util.List;
import java.util.function.Function;

// @see https://blog.csdn.net/qq_42093488/article/details/124956848 批处理
public abstract class DynamicSqlSessionBatch extends DynamicSqlSessionPreparedStatement {

    protected Integer batchSize = 500;

    /**
     * 批量执行
     **/
    public Long executeBatch(DynamicMappedStatement dms, Object parameter) {
        int myBatchSize = dms.getBatchSize() > 0 ? dms.getBatchSize() : this.batchSize;
        return executeBatch(dms, parameter, myBatchSize);
    }

    /**
     * 批量执行
     *
     * @param parameter ignore
     * @param parameter parameter 类型必须是 Collection
     * @param batchSize ignore
     * @return ignore
     */
    public long executeBatch(DynamicMappedStatement dms, Object parameter, int batchSize) {
        batchSize = batchSize < 1 ? this.batchSize : batchSize;
        if (!(parameter instanceof Collection)) {
            throw new RuntimeException(" parameter type not is  Collection.", null);
        }
        return executeBatch(dms, (Collection) parameter, batchSize);
    }

    /**
     * 执行批量操作
     *
     * @param list      数据集合
     * @param batchSize 批次大小
     * @return 操作结果
     * @since 3.4.0
     */
    protected <E> long executeBatch(DynamicMappedStatement dms, Collection<E> list, int batchSize) {
        Assert.isFalse(batchSize < 1, "batchSize must not be less than one");
        if (CollectionUtils.isEmpty(list)) {
            return 0;
        }
        return executeBatch(sqlSession -> {
            int size = list.size();
            int idxLimit = Math.min(batchSize, size);
            int i = 1;
            int total = 0;
            for (E element : list) {
                getDynamicSqlSession().execute(sqlSession, dms, element, null, true);
                if (i == idxLimit) {
                    List<BatchResult> batchResults = sqlSession.flushStatements();
                    idxLimit = Math.min(idxLimit + batchSize, size);
                    if (batchResults != null) {
                        total += batchResults.stream().mapToInt(b -> Arrays.stream(b.getUpdateCounts()).sum()).sum();
                    }
                }
                i++;
            }
            return total;
        });
    }

    /**
     * 执行批量操作
     *
     * @param function function
     * @return 操作结果
     * @since 3.4.0
     */
    protected long executeBatch(Function<SqlSession, Integer> function) {
        SqlSession sqlSession = getSqlSqlSessionFactory().openSession(ExecutorType.BATCH);
        try {
            long count = function.apply(sqlSession);
            // 非事务环境下强制commit，事务情况下该commit相当于无效
            //  sqlSession.commit(!TransactionSynchronizationManager.isActualTransactionActive());
            return count;
        } catch (Throwable t) {
            sqlSession.rollback();
            throw t;
        } finally {
            sqlSession.close();
        }
    }
}
